# AMAEventMakerHelper.py
# 
# Copyright (C) 2010 Hurng-Chun Lee <hurngchunlee@gmail.com>
# 
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

from Ganga.Utility.logging import getLogger

def patchPPTask(gtId, gtId_pp, trfIdList=[], searchAlt=True):
    '''
    This methods provides the possibility to patch some of the transforms of the AMAEvent generating post-processing tasks.

    By setting 'searchAlt' to True, it searches into the origianl GangaTask of AMAEvent generation to get another location
    of running the particular transform (as GangaTask may generate same AMAEvent files at 2 different locations.)

    @param gtId specifies GangaTask id of the AMAEvent generating task
    @param gtId_pp specifies GangaTask id of the AMAEvent generating post-processing task
    @param trfIdList specifies Transforms id list of the AMAEvent generating post-processing task that are going to be patched
    @param searchAlt searches possible alternative location for running the post-process
    '''

    logger = getLogger('AutoD3PDMakerUserInterface')

    if not trfIdList:
        id = 0
        for trf in tasks(gtId_pp).transforms:
            if trf.status not in ['completed']:
                trfIdList.append( id )
            id += 1

    ## pause the running transforms
    for id in trfIdList:
        if tasks(gtId_pp).transforms[id].status in ['running']:
            ## pause the running transforms
            logger.debug('pausing transform %d' % id)
            tasks(gtId_pp).transforms[id].pause()

        ## removing old jobs
        logger.debug('removing old jobs ...')
        for sj in tasks(gtId_pp).transforms[id].getJobs():
            sj.remove()

        if searchAlt:
            trf_app    = tasks(gtId_pp).transforms[id].application.copy()

            input_ds   = trf_app.args[0]
            input_f    = trf_app.args[1]
            output_loc = trf_app.args[3]

            pp_out_info = [ input_ds, input_f, output_loc ]

            ## try to resolve from the original GangaTask with alternative locatioons
            altList = []

            for gtTrf in tasks(gtId).transforms:
                for partId in range(1, gtTrf.n_all()+1):

                    outList = []

                    for sj in gtTrf.getPartitionJobs(partId):
                        if sj.status == 'completed':
                            for out in sj.outputdata.output:
                                out_info = out.split(',')
                                outList.append( [out_info[0],out_info[1],out_info[-1]] )

                    if pp_out_info in outList:
                        altList = outList
                        altList.remove( pp_out_info )
                        break

                ## if found the altList already, break the look immediately
                if altList:
                    break

            ## althernative does exist, reset the transform
            if altList:
                
                altInfo = altList[0]
                logger.debug('trying location %s for post-processing file %s:%s' % (atlInfo[-1], atlInfo[0], altInfo[1]) )

                trf_app.args[0] = altInfo[0]
                trf_app.args[1] = altInfo[1]
                trf_app.args[3] = altInfo[2]

                new_trf = Transform()
                new_trf.application = trf_app
                new_trf.backend = LCG()
                new_trf.backend.requirements.sites = [ atlInfo[-1] ] 
                new_trf.setPartitionStatus(1,'ready')

                tasks(gtId_pp).transforms[id] = new_trf

            else:
                logger.warning('no alternative location of %s:%s' % (input_ds, input_f) )


        ## rerun the transform to run
        logger.debug('rerun the transform: t_%d:trf_%d' % (gtId_pp, id))
        tasks(gtId_pp).transforms[id].run()

    return True
